! ==============================================================================
! This file is part of CUPP (Community Unified Physics Package) since 2023.
! ==============================================================================

module mp_lin_mod

  use cupp_math_mod
  use cupp_mesh_mod
  use atm_types_mod
  use mp_const_mod
  use tracer_mod

  implicit none

  private

  public mp_lin_init
  public mp_lin_final
  public mp_lin_run

  ! Equations for the water substances:
  !
  !   d q
  !   ---  = - adv(q) + dif(q) - prod_r - prod_s - prod_g
  !   d t
  !
  !   d lr
  !   ---- = - adv(lr) + dif(lr) + prod_r + fallout_r
  !   d t
  !
  !   d ls
  !   ---- = - adv(ls) + dif(ls) + prod_s + fallout_s
  !   d t
  !
  !   d lg
  !   ---- = - adv(lg) + dif(lg) + prod_g + fallout_g
  !   d t

  ! Intercept parameters of the raindrop size distribution.
  real(r8), parameter :: n0r          = 8.0e6_r8         ! Rain intercept parameter (m-4)
  real(r8), parameter :: n0s          = 3.0e6_r8         ! Snow intercept parameter (m-4)
  real(r8), parameter :: n0g          = 4.0e4_r8         ! Graupel intercept parameter (m-4)
  real(r8), parameter :: qic0         = 7.0e-4_r8        ! Threshold for aggregation to occur (kg kg-1)
  real(r8), parameter :: rho0         = 1.29_r8          ! Surface air density (kg m-3)
  real(r8), parameter :: vta          = 2115.0e-2_r8     ! Constant a in empirical formula for terminal velocity of rain (m**(1-b) s-1)
  real(r8), parameter :: vtb          = 0.8_r8           ! Constant b in empirical formula for terminal velocity of rain
  real(r8), parameter :: vtc          = 152.93e-2_r8     ! Constant c in empirical formula for terminal velocity of snow (m**(1-d) s-1)
  real(r8), parameter :: vtd          = 0.25_r8          ! Constant d in empirical formula for terminal velocity of snow
  real(r8), parameter :: cdrag        = 0.6_r8           ! Drag coefficient for hailstone
  real(r8), parameter :: gamma_4p5    = gamma(4.5_r8)    ! Gamma function of 4.5
  real(r8), parameter :: gamma_vtd_p3 = gamma(vtd + 3)
  real(r8), parameter :: gamma_vtd_p4 = gamma(vtd + 4)
  real(r8), parameter :: gamma_vtb_p4 = gamma(vtb + 4)

  type lin_data_type
    integer :: id
    type(cupp_mesh_type), pointer :: mesh => null()
    real(r8), allocatable, dimension(:,:) :: qsi      ! Saturation specific humidity over ice (kg kg-1)
    real(r8), allocatable, dimension(:,:) :: qsw      ! Saturation specific humidity over water (kg kg-1)
    real(r8), allocatable, dimension(:,:) :: sqrt_rho ! sqrt(rho0 / rho)
    real(r8), allocatable, dimension(:,:) :: vt_r     ! Terminal velocity of rain (m s-1)
    real(r8), allocatable, dimension(:,:) :: vt_s     ! Terminal velocity of snow (m s-1)
    real(r8), allocatable, dimension(:,:) :: vt_g     ! Terminal velocity of graupel (m s-1)
    real(r8), allocatable, dimension(:,:) :: lambda_r ! Slope parameter of rain size distribution
    real(r8), allocatable, dimension(:,:) :: lambda_s ! Slope parameter of snow size distribution
    real(r8), allocatable, dimension(:,:) :: lambda_g ! Slope parameter of graupel size distribution
    ! real(r8), allocatable, dimension(:,:) ::
    real(r8), allocatable, dimension(:,:) :: psaut
    real(r8), allocatable, dimension(:,:) :: psaci
  contains
    procedure :: init  => lin_data_init
    procedure :: clear => lin_data_clear
    final :: lin_data_final
  end type lin_data_type

  type(lin_data_type), allocatable :: lin_data(:)

contains

  subroutine mp_lin_init(namelist, mesh)

    character(*), intent(in) :: namelist
    type(cupp_mesh_type), intent(in) :: mesh(:)

    integer iblk

    call mp_lin_final()

    ! Add tracers.
    call tracer_add('qv', 'Water vapor', 'kg kg-1')
    ! - Cloud particles
    call tracer_add('ql', 'Cloud water', 'kg kg-1')
    call tracer_add('qi', 'Cloud ice'  , 'kg kg-1')
    ! - Precipitation particles
    call tracer_add('qs', 'Snow'       , 'kg kg-1')
    call tracer_add('qr', 'Rain'       , 'kg kg-1')
    call tracer_add('qg', 'Graupel'    , 'kg kg-1')

    allocate(lin_data(size(mesh)))
    do iblk = 1, size(mesh)
      call lin_data(iblk)%init(mesh(iblk))
    end do

  end subroutine mp_lin_init

  subroutine mp_lin_final()

    if (allocated(lin_data)) deallocate(lin_data)

  end subroutine mp_lin_final

  subroutine mp_lin_run(state, tend, dt_phys, dt_mp)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_mp

    integer i, k
    real(r8) es

    ! Calculate some variables needed for the microphysics.
    associate (mesh     => state%mesh, &
               t        => state%t   , &
               p        => state%p   , &
               rho      => state%rho , &
               sqrt_rho => lin_data(state%mesh%id)%sqrt_rho, &
               qsi      => lin_data(state%mesh%id)%qsi, &
               qsw      => lin_data(state%mesh%id)%qsw)
    do k = 1, mesh%nlev
      do i = 1, mesh%ncol_with_halo
        es = 611.2_r8 * exp(17.67_r8 * (t(i,k) - t0) / (t(i,k) - 29.65_r8))
        qsw(i,k) = rd_o_rv * es / (p(i,k) - es)
        if (t(i,k) < t0) then
          es = 611.2_r8 * exp(21.8745584_r8 * (t(i,k) - t00) / (t(i,k) - 7.66_r8))
          qsi(i,k) = rd_o_rv * es / (p(i,k) - es)
          if (t(i,k) < t0 - 40) qsw(i,k) = qsi(i,k)
        else
          qsi(i,k) = qsw(i,k)
        end if
        sqrt_rho(i,k) = sqrt(rho0 / rho(i,k))
      end do
    end do
    end associate

    call terminal_velocity(state)
    call viscosity(state)

  end subroutine mp_lin_run

  subroutine terminal_velocity(state)

    type(atm_state_type), intent(in) :: state

    integer i, k
    real(r8) qt

    associate (mesh     => state%mesh                      , &
               t        => state%t                         , & ! in
               p        => state%p                         , & ! in
               rho      => state%rho                       , & ! in
               sqrt_rho => lin_data(state%mesh%id)%sqrt_rho, & ! in
               qv       => state%qv                        , & ! in
               ql       => state%qc                        , & ! in
               qi       => state%qi                        , & ! in
               qr       => state%qr                        , & ! in
               qs       => state%qs                        , & ! in
               qg       => state%qg                        , & ! in
               qsi      => lin_data(state%mesh%id)%qsi     , & ! in
               vt_r     => lin_data(state%mesh%id)%vt_r    , & ! out
               vt_s     => lin_data(state%mesh%id)%vt_s    , & ! out
               vt_g     => lin_data(state%mesh%id)%vt_g    , & ! out
               lambda_r => lin_data(state%mesh%id)%lambda_r, & ! out
               lambda_s => lin_data(state%mesh%id)%lambda_s, & ! out
               lambda_g => lin_data(state%mesh%id)%lambda_g)   ! out
    do k = 1, mesh%nlev
      do i = 1, mesh%ncol_with_halo
        qt = ql(i,k) + qi(i,k) + qr(i,k) + qs(i,k) + qg(i,k)
        if (.not. (qv(i,k) + ql(i,k) + qi(i,k) < qsi(i,k) .and. qt == 0)) then
          if (qr(i,k) > 1.0e-8_r8) then
            lambda_r(i,k) = sqrt(sqrt(pi * rhow * n0r / (rho(i,k) * qr(i,k))))
            vt_r    (i,k) = vta * gamma_vtb_p4 * sqrt_rho(i,k) / (6 * lambda_r(i,k)**vtb)
          else
            lambda_r(i,k) = 0
            vt_r    (i,k) = 0
          end if
          if (qs(i,k) > 1.0e-8_r8) then
            lambda_s(i,k) = sqrt(sqrt(pi * rhos * n0s / (rho(i,k) * qs(i,k))))
            vt_s    (i,k) = vtc * gamma_vtd_p4 * sqrt_rho(i,k) / (6 * lambda_s(i,k)**vtd)
          else
            lambda_s(i,k) = 0
            vt_s    (i,k) = 0
          end if
          if (qg(i,k) > 1.0e-8_r8) then
            lambda_g(i,k) = sqrt(sqrt(pi * rhog * n0g / (rho(i,k) * qg(i,k))))
            vt_g    (i,k) = gamma_4p5 * sqrt(4 * g * rhog / (3 * cdrag * rho(i,k))) / (6 * sqrt(lambda_g(i,k)))
          else
            lambda_g(i,k) = 0
            vt_g    (i,k) = 0
          end if
        end if
      end do
    end do
    end associate

  end subroutine terminal_velocity

  subroutine viscosity(state)

    type(atm_state_type), intent(in) :: state

    integer i, k


  end subroutine viscosity

  subroutine snow_production(state, tend, dt_phys, dt_mp)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_mp

    ! Snow production for T < 0 degC
    ! - (psaut) Autoconversion of cloud ice to snow or ice crystal aggregation
    call snow_aggregation(state, tend, dt_phys, dt_mp)
    ! - (psaci) Accretion of cloud ice by snow
    call snow_accretion(state, tend, dt_phys, dt_mp)
    ! - (Ps_acw) Accretion of cloud water by snow
    ! - (Ps_fw ) Bergeron process (deposition and riming) - transfer of cloud water to snow
    ! - (Ps_fi ) Transfer rate of cloud ice to snow through growth of Bergeron process embryos
    ! - (Pr_aci)
    ! - (Pi_acr)
    ! - (Pg_acs)
    ! - (Pg_aut)
    ! - (Pr_acs)
    ! - (Ps_acr)
    ! - (Ps_sur)
    ! - (Ps_dep)

    ! Snow production for T >= 0 degC
    ! - (Ps_mlt)
    ! - (Pg_acs)

  end subroutine snow_production

  subroutine snow_aggregation(state, tend, dt_phys, dt_mp)

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_mp

    real(r8) a1   ! Aggregation coefficient (s-1)
    real(r8) qic  ! Threshold for aggregation to occur (kg kg-1)
    integer i, k

    associate (mesh  => state%mesh                   , &
               t     => state%t                      , & ! in
               rho   => state%rho                    , & ! in
               qi    => state%qi                     , & ! in
               psaut => lin_data(state%mesh%id)%psaut)   ! out
    do k = 1, mesh%nlev
      do i = 1, mesh%ncol_with_halo
        if (t(i,k) < t0) then
          a1 = 1.0e-3_r8 * exp(0.025_r8 * (t(i,k) - t0))
          if (t(i,k) < t0 - 20) then
            qic = 1.0e-3_r8 / rho(i,k) * exp(     &
              -7.6_r8 + 4 * exp(-0.2443e-3_r8 * ( &
                abs(t(i,k) - t0) - 20             &
              )**2.455_r8))
          else
            qic = qic0
          end if
          psaut(i,k) = max(0.0_r8, (1 - exp(-a1 / dt_mp)) * (qi(i,k) - qic) / dt_mp)
        end if
      end do
    end do
    end associate

  end subroutine snow_aggregation

  subroutine snow_accretion(state, tend, dt_phys, dt_mp)

    ! Geometric sweep-out concept integrated over all snow sizes for the assumed snow size distribution.

    type(atm_state_type), intent(in) :: state
    type(atm_tend_type), intent(inout) :: tend
    real(r8), intent(in) :: dt_phys
    real(r8), intent(in) :: dt_mp

    real(r8) esi ! Snow collection efficiency
    real(r8) tmp1, tmp2
    integer i, k

    associate (mesh     => state%mesh                      , &
               t        => state%t                         , & ! in
               qi       => state%qi                        , & ! in
               sqrt_rho => lin_data(state%mesh%id)%sqrt_rho, & ! in
               lambda_s => lin_data(state%mesh%id)%lambda_s, & ! in
               psaci    => lin_data(state%mesh%id)%psaci   )   ! out
    do k = 1, mesh%nlev
      do i = 1, mesh%ncol_with_halo
        ! The accretion of cloud ice by snow is an aggregation process
        ! which occurs if the temperature is below 0 degC.
        if (t(i,k) < t0) then
          esi = exp(0.025_r8 * (t(i,k) - t0))
          tmp1 = pio4 * n0s * vtc * gamma_vtd_p3 * sqrt_rho(i,k) / lambda_s(i,k)**(3 + vtd)
          tmp2 = tmp1 * esi
          psaci(i,k) = (1 - exp(-tmp2 * dt_mp)) * qi(i,k) / dt_mp
        end if
      end do
    end do
    end associate

  end subroutine snow_accretion

  subroutine lin_data_init(this, mesh)

    class(lin_data_type), intent(inout) :: this
    type(cupp_mesh_type), intent(in), target :: mesh

    call this%clear()

    this%id   = mesh%id
    this%mesh => mesh

    allocate(this%qsi     (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%qsw     (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%sqrt_rho(mesh%ncol_with_halo,mesh%nlev))
    allocate(this%vt_r    (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%vt_s    (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%lambda_r(mesh%ncol_with_halo,mesh%nlev))
    allocate(this%lambda_s(mesh%ncol_with_halo,mesh%nlev))
    allocate(this%lambda_g(mesh%ncol_with_halo,mesh%nlev))
    allocate(this%psaut  (mesh%ncol_with_halo,mesh%nlev))
    allocate(this%psaci  (mesh%ncol_with_halo,mesh%nlev))

  end subroutine lin_data_init

  subroutine lin_data_clear(this)

    class(lin_data_type), intent(inout) :: this

    this%id   = 0
    this%mesh => null()

    if (allocated(this%qsi     )) deallocate(this%qsi     )
    if (allocated(this%qsw     )) deallocate(this%qsw     )
    if (allocated(this%sqrt_rho)) deallocate(this%sqrt_rho)
    if (allocated(this%vt_r    )) deallocate(this%vt_r    )
    if (allocated(this%vt_s    )) deallocate(this%vt_s    )
    if (allocated(this%lambda_r)) deallocate(this%lambda_r)
    if (allocated(this%lambda_s)) deallocate(this%lambda_s)
    if (allocated(this%lambda_g)) deallocate(this%lambda_g)
    if (allocated(this%psaut  )) deallocate(this%psaut  )
    if (allocated(this%psaci  )) deallocate(this%psaci  )

  end subroutine lin_data_clear

  subroutine lin_data_final(this)

    type(lin_data_type), intent(inout) :: this

    call this%clear()

  end subroutine lin_data_final

end module mp_lin_mod